from typing import Type, Callable, Tuple, Set, List, Union

import torch
import torch.nn as nn
from timm.models.efficientnet_blocks import SqueezeExcite, DepthwiseSeparableConv
from timm.models.layers import drop_path, trunc_normal_, Mlp, DropPath


def _gelu_ignore_parameters(
        *args,
        **kwargs
) -> nn.Module:
    """ 一个小技巧，用于在 Timm 的 DepthwiseSeparableConv 中忽略 inplace=True 参数。

    参数:
        *args: 忽略。
        **kwargs: 忽略。

    返回:
        activation (nn.Module): GELU 激活函数。
    """
    activation = nn.GELU()
    return activation


class MBConv(nn.Module):
    """ MBConv 模块，参考论文：https://arxiv.org/pdf/2204.01697.pdf。

        无下采样：
        x ← x + Proj(SE(DWConv(Conv(Norm(x)))))

        有下采样：
        x ← Proj(Pool2D(x)) + Proj(SE(DWConv ↓(Conv(Norm(x))))).

        Conv 是一个 1x1 卷积，后接批归一化层和 GELU 激活。
        SE 是 Squeeze-Excitation 层。
        Proj 是缩减通道的 1x1 卷积。

        注意：此实现与原始 MobileNet 实现略有不同！

    参数:
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        act_layer (Type[nn.Module], optional): 使用的激活层类型。默认值：nn.GELU
        norm_layer (Type[nn.Module], optional): 使用的归一化层类型。默认值：nn.BatchNorm2d
        drop_path (float, optional): 训练期间应用的 Dropout 率。默认值 0。
    """

    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            drop_path: float = 0.,
            debug=False
    ) -> None:
        """ 构造函数 """
        # 调用父类构造函数
        super(MBConv, self).__init__()
        self.module_name = '[MBConv]'
        self.debug = debug
        # 保存参数
        self.drop_path_rate: float = drop_path
        # 如果使用 GELU，则忽略 inplace 参数
        if act_layer == nn.GELU:
            act_layer = _gelu_ignore_parameters
        # 构建主路径
        self.main_path = nn.Sequential(
            norm_layer(in_channels),
            nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=(1, 1)),
            DepthwiseSeparableConv(in_chs=in_channels, out_chs=out_channels, stride=1,
                                   act_layer=act_layer, norm_layer=norm_layer, drop_path_rate=drop_path),
            SqueezeExcite(in_chs=out_channels, rd_ratio=0.25),
            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 1))
        )
        # 构建跳跃连接
        self.skip_path = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, 1)) \
            if (in_channels != out_channels) else nn.Identity()

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """ 前向传播

        参数:
            input (torch.Tensor): 输入张量，形状为 [B, C_in, H, W]。

        返回:
            output (torch.Tensor): 输出张量，形状为 [B, C_out, H (// 2), W (// 2)]（下采样是可选的）。
        """
        output = self.main_path(input)
        if self.drop_path_rate > 0.:
            output = drop_path(output, self.drop_path_rate, self.training)
        output = output + self.skip_path(input)
        return output


def window_partition(
        input: torch.Tensor,
        window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """ 窗口分割函数

    参数:
        input (torch.Tensor): 输入张量，形状为 [B, C, H, W]。
        window_size (Tuple[int, int], optional): 窗口大小。默认值 (7, 7)

    返回:
        windows (torch.Tensor): 分割后的张量，形状为 [B * windows, window_size[0], window_size[1], C]。
    """
    # 获取输入的尺寸
    B, C, H, W = input.shape
    # 分割输入
    windows = input.view(B, C, H // window_size[0], window_size[0], W // window_size[1], window_size[1])
    # 置换和重塑为 [B * windows, window_size[0], window_size[1], channels]
    windows = windows.permute(0, 2, 4, 3, 5, 1).contiguous().view(-1, window_size[0], window_size[1], C)
    return windows


def window_reverse(
        windows: torch.Tensor,
        original_size: Tuple[int, int],
        window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """ 恢复窗口分割

    参数:
        windows (torch.Tensor): 窗口张量，形状为 [B * windows, window_size[0], window_size[1], C]。
        original_size (Tuple[int, int]): 原始尺寸。
        window_size (Tuple[int, int], optional): 窗口大小。默认值 (7, 7)

    返回:
        output (torch.Tensor): 恢复后的张量，形状为 [B, C, original_size[0], original_size[1]]。
    """
    # 获取高度和宽度
    H, W = original_size
    # 计算批大小
    B = int(windows.shape[0] / (H * W / window_size[0] / window_size[1]))
    # 恢复张量
    output = windows.view(B, H // window_size[0], W // window_size[1], window_size[0], window_size[1], -1)
    output = output.permute(0, 5, 1, 3, 2, 4).contiguous().view(B, -1, H, W)
    return output


def grid_partition(
        input: torch.Tensor,
        grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """ 网格分割函数

    参数:
        input (torch.Tensor): 输入张量，形状为 [B, C, H, W]。
        grid_size (Tuple[int, int], optional): 网格大小。默认值 (7, 7)

    返回:
        grid (torch.Tensor): 分割后的张量，形状为 [B * grids, grid_size[0], grid_size[1], C]。
    """
    # 获取输入的尺寸
    B, C, H, W = input.shape
    # 分割输入
    grid = input.view(B, C, grid_size[0], H // grid_size[0], grid_size[1], W // grid_size[1])
    # 置换和重塑
    grid = grid.permute(0, 3, 5, 2, 4, 1).contiguous().view(-1, grid_size[0], grid_size[1], C)
    return grid


def grid_reverse(
        grid: torch.Tensor,
        original_size: Tuple[int, int],
        grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """ 恢复网格分割

    参数:
        grid (torch.Tensor): 网格张量，形状为 [B * grids, grid_size[0], grid_size[1], C]。
        original_size (Tuple[int, int]): 原始尺寸。
        grid_size (Tuple[int, int], optional): 网格大小。默认值 (7, 7)

    返回:
        output (torch.Tensor): 恢复后的张量，形状为 [B, C, original_size[0], original_size[1]]。
    """
    # 获取高度、宽度和通道数
    (H, W), C = original_size, grid.shape[-1]
    # 计算批大小
    B = int(grid.shape[0] / (H * W / grid_size[0] / grid_size[1]))
    # 恢复张量
    output = grid.view(B, H // grid_size[0], W // grid_size[1], grid_size[0], grid_size[1], C)
    output = output.permute(0, 5, 3, 1, 4, 2).contiguous().view(B, C, H, W)
    return output


def get_relative_position_index(
        win_h: int,
        win_w: int
) -> torch.Tensor:
    """ 为窗口内的每个标记生成成对的相对位置索引。
        来源于 Timm 的 Swin V1 实现。

    参数:
        win_h (int): 窗口/网格高度。
        win_w (int): 窗口/网格宽度。

    返回:
        relative_coords (torch.Tensor): 成对的相对位置索引 [height * width, height * width]。
    """
    coords = torch.stack(torch.meshgrid([torch.arange(win_h), torch.arange(win_w)], indexing='ij'))
    coords_flatten = torch.flatten(coords, 1)
    relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
    relative_coords = relative_coords.permute(1, 2, 0).contiguous()
    relative_coords[:, :, 0] += win_h - 1
    relative_coords[:, :, 1] += win_w - 1
    relative_coords[:, :, 0] *= 2 * win_w - 1
    return relative_coords.sum(-1)


class RelativeSelfAttention(nn.Module):
    """ 类似 Swin V1 的相对自注意力。实现参考 Timm 的 Swin V1 实现。

    参数:
        in_channels (int): 输入通道数。
        num_heads (int, optional): 注意力头的数量。默认值 32
        grid_window_size (Tuple[int, int], optional): 网格/窗口大小。默认值 (7, 7)
        attn_drop (float, optional): 注意力权重的 Dropout 比例。默认值 0.0
        drop (float, optional): 输出的 Dropout 比例。默认值 0.0
    """

    def __init__(
            self,
            in_channels: int,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.,
            drop: float = 0.
    ) -> None:
        """ 构造函数 """
        # 调用父类构造函数
        super(RelativeSelfAttention, self).__init__()
        # 保存参数
        self.in_channels: int = in_channels
        self.num_heads: int = num_heads
        self.grid_window_size: Tuple[int, int] = grid_window_size
        self.scale: float = num_heads ** -0.5
        self.attn_area: int = grid_window_size[0] * grid_window_size[1]
        # 初始化层
        self.qkv_mapping = nn.Linear(in_features=in_channels, out_features=3 * in_channels, bias=True)
        self.attn_drop = nn.Dropout(p=attn_drop)
        self.proj = nn.Linear(in_features=in_channels, out_features=in_channels, bias=True)
        self.proj_drop = nn.Dropout(p=drop)
        self.softmax = nn.Softmax(dim=-1)
        # 定义相对位置偏置参数表，形状：2*Wh-1 * 2*Ww-1, nH
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * grid_window_size[0] - 1) * (2 * grid_window_size[1] - 1), num_heads))

        # 获取窗口内每个标记的成对相对位置索引
        self.register_buffer("relative_position_index", get_relative_position_index(grid_window_size[0],
                                                                                    grid_window_size[1]))
        # 初始化相对位置偏置
        trunc_normal_(self.relative_position_bias_table, std=.02)

    def _get_relative_positional_bias(
            self
    ) -> torch.Tensor:
        """ 返回相对位置偏置。

        返回:
            relative_position_bias (torch.Tensor): 相对位置偏置。
        """
        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)].view(self.attn_area, self.attn_area, -1)
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
        return relative_position_bias.unsqueeze(0)

    def forward(
            self,
            input: torch.Tensor
    ) -> torch.Tensor:
        """ 前向传播。

        参数:
            input (torch.Tensor): 输入张量，形状为 [B_, N, C]。

        返回:
            output (torch.Tensor): 输出张量，形状为 [B_, N, C]。
        """
        # 获取输入的形状
        B_, N, C = input.shape
        # 执行查询、键、值映射
        qkv = self.qkv_mapping(input).reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        # 缩放查询
        q = q * self.scale
        # 计算注意力图
        attn = self.softmax(q @ k.transpose(-2, -1) + self._get_relative_positional_bias())
        # 映射值
        output = (attn @ v).transpose(1, 2).reshape(B_, N, -1)
        # 最终投影和 Dropout
        output = self.proj(output)
        output = self.proj_drop(output)
        return output


class MaxViTTransformerBlock(nn.Module):
    """ MaxViT Transformer 块。

        使用块分割：
        x ← x + Unblock(RelAttention(Block(LN(x))))
        x ← x + MLP(LN(x))

        使用网格分割：
        x ← x + Ungrid(RelAttention(Grid(LN(x))))
        x ← x + MLP(LN(x))

        层归一化 (LN) 在网格/窗口分割后应用，以避免多次重塑操作。
        最终输出的逆操作也是出于同样的原因。

    参数:
        in_channels (int): 输入通道数。
        partition_function (Callable): 使用的分割函数（网格或窗口分割）。
        reverse_function (Callable): 使用的逆函数（网格或窗口逆操作）。
        num_heads (int, optional): 注意力头的数量。默认值 32
        grid_window_size (Tuple[int, int], optional): 网格/窗口大小。默认值 (7, 7)
        attn_drop (float, optional): 注意力权重的 Dropout 比例。默认值 0.0
        drop (float, optional): 输出的 Dropout 比例。默认值 0.0
        drop_path (float, optional): 路径的 Dropout 比例。默认值 0.0
        mlp_ratio (float, optional): MLP 隐藏层维度与嵌入维度的比例。默认值 4.0
        act_layer (Type[nn.Module], optional): 使用的激活层类型。默认值 nn.GELU
        norm_layer (Type[nn.Module], optional): 使用的归一化层类型。默认值 nn.LayerNorm
    """

    def __init__(
            self,
            in_channels: int,
            partition_function: Callable,
            reverse_function: Callable,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.,
            drop: float = 0.,
            drop_path: float = 0.,
            mlp_ratio: float = 4.,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            debug=False
    ) -> None:
        """ 构造函数 """
        super(MaxViTTransformerBlock, self).__init__()
        self.module_name = '[MaxViTTransformerBlock]'
        self.debug = debug
        # 保存参数
        self.partition_function: Callable = partition_function
        self.reverse_function: Callable = reverse_function
        self.grid_window_size: Tuple[int, int] = grid_window_size
        # 初始化层
        self.norm_1 = norm_layer(in_channels)
        self.attention = RelativeSelfAttention(
            in_channels=in_channels,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm_2 = norm_layer(in_channels)
        self.mlp = Mlp(
            in_features=in_channels,
            hidden_features=int(mlp_ratio * in_channels),
            act_layer=act_layer,
            drop=drop
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """ 前向传播。

        参数:
            input (torch.Tensor): 输入张量，形状为 [B, C_in, H, W]。

        返回:
            output (torch.Tensor): 输出张量，形状为 [B, C_out, H (// 2), W (// 2)]。
        """
        # 保存原始形状
        B, C, H, W = input.shape
        # 分割
        input_partitioned = self.partition_function(input, self.grid_window_size)
        input_partitioned = input_partitioned.view(-1, self.grid_window_size[0] * self.grid_window_size[1], C)
        # 归一化、注意力、Dropout
        output = input_partitioned + self.drop_path(self.attention(self.norm_1(input_partitioned)))
        # 归一化、MLP、Dropout
        output = output + self.drop_path(self.mlp(self.norm_2(output)))
        # 恢复分割
        output = self.reverse_function(output, (H, W), self.grid_window_size)
        return output


class MaxViTBlock(nn.Module):
    """ MaxViT 块，由 MBConv 块、块注意力和网格注意力组成。

    参数:
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        num_heads (int, optional): 注意力头的数量。默认值 32
        grid_window_size (Tuple[int, int], optional): 网格/窗口大小。默认值 (7, 7)
        attn_drop (float, optional): 注意力权重的 Dropout 比例。默认值 0.0
        drop (float, optional): 输出的 Dropout 比例。默认值 0.0
        drop_path (float, optional): 路径的 Dropout 比例。默认值 0.0
        mlp_ratio (float, optional): MLP 隐藏层维度与嵌入维度的比例。默认值 4.0
        act_layer (Type[nn.Module], optional): 使用的激活层类型。默认值 nn.GELU
        norm_layer (Type[nn.Module], optional): 使用的归一化层类型。默认值 nn.BatchNorm2d
        norm_layer_transformer (Type[nn.Module], optional): Transformer 中的归一化层。默认值 nn.LayerNorm
    """

    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.,
            drop: float = 0.,
            drop_path: float = 0.,
            mlp_ratio: float = 4.,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            norm_layer_transformer: Type[nn.Module] = nn.LayerNorm,
            debug=False,
    ) -> None:
        """ 构造函数 """
        # 调用父类构造函数
        super(MaxViTBlock, self).__init__()
        self.module_name = '[MaxViTBlock]'
        self.debug = debug
        if self.debug: print(self.module_name, '[in_channels]', in_channels)
        if self.debug: print(self.module_name, '[out_channels]', out_channels)
        if self.debug: print(self.module_name, '[num_heads]', num_heads)
        if self.debug: print(self.module_name, '[grid_window_size]', grid_window_size)
        if self.debug: print(self.module_name, '[attn_drop]', attn_drop)
        if self.debug: print(self.module_name, '[drop]', drop)
        if self.debug: print(self.module_name, '[drop_path]', drop_path)
        if self.debug: print(self.module_name, '[mlp_ratio]', mlp_ratio)
        # 初始化 MBConv 块
        self.mb_conv = MBConv(
            in_channels=in_channels,
            out_channels=out_channels,
            act_layer=act_layer,
            norm_layer=norm_layer,
            drop_path=drop_path,
            debug=debug,
        )
        # 初始化块和网格 Transformer
        self.block_transformer = MaxViTTransformerBlock(
            in_channels=out_channels,
            partition_function=window_partition,
            reverse_function=window_reverse,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop,
            drop_path=drop_path,
            mlp_ratio=mlp_ratio,
            act_layer=act_layer,
            norm_layer=norm_layer_transformer,
            debug=debug,
        )
        self.grid_transformer = MaxViTTransformerBlock(
            in_channels=out_channels,
            partition_function=grid_partition,
            reverse_function=grid_reverse,
            num_heads=num_heads,
            grid_window_size=grid_window_size,
            attn_drop=attn_drop,
            drop=drop,
            drop_path=drop_path,
            mlp_ratio=mlp_ratio,
            act_layer=act_layer,
            norm_layer=norm_layer_transformer,
            debug=debug,
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """ 前向传播。

        参数:
            input (torch.Tensor): 输入张量，形状为 [B, C_in, H, W]

        返回:
            output (torch.Tensor): 输出张量，形状为 [B, C_out, H // 2, W // 2]（下采样是可选的）
        """
        output = self.mb_conv(input)
        output = self.block_transformer(output)
        output = self.grid_transformer(output)
        return output


class DeconvModule(nn.Module):
    """UNet 解码器中使用反卷积上采样模块（2 倍上采样）。

    该模块在 UNet 的解码器中使用反卷积对特征图进行上采样。

    参数:
        in_channels (int): 输入通道数。
        out_channels (int): 输出通道数。
        norm_layer: 使用的归一化层。
            默认值：nn.BatchNorm2d。
        act_layer (dict | None): 使用的激活层。
            默认值：nn.Mish。
        kernel_size (int): 卷积层的核大小。默认值：4。
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 norm_layer=nn.BatchNorm2d,
                 act_layer=nn.Mish,
                 kernel_size=4,
                 scale_factor=2):
        super(DeconvModule, self).__init__()

        assert (kernel_size - scale_factor >= 0) and \
               (kernel_size - scale_factor) % 2 == 0, \
            f'kernel_size 应大于等于 scale_factor 且 (kernel_size - scale_factor) 应为偶数，' \
            f'当前 kernel size 为 {kernel_size}，scale factor 为 {scale_factor}。'

        stride = scale_factor
        padding = (kernel_size - scale_factor) // 2
        deconv = nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding)

        norm = norm_layer(out_channels)
        activate = act_layer()
        self.deconv_upsamping = nn.Sequential(deconv, norm, activate)

    def forward(self, x):
        """前向函数。"""
        out = self.deconv_upsamping(x)
        return out


class MaxViTStage(nn.Module):
    """ MaxViT 的阶段。

    参数:
        depth (int): 阶段的深度。
        in_channels (int): 输入通道数。
        skip_channels (int): 跳跃连接的通道数。
        out_channels (int): 输出通道数。
        num_heads (int, optional): 注意力头的数量。默认值 32
        grid_window_size (Tuple[int, int], optional): 网格/窗口大小。默认值 (7, 7)
        attn_drop (float, optional): 注意力权重的 Dropout 比例。默认值 0.0
        drop (float, optional): 输出的 Dropout 比例。默认值 0.0
        drop_path (float, optional): 路径的 Dropout 比例。默认值 0.0
        mlp_ratio (float, optional): MLP 隐藏层维度与嵌入维度的比例。默认值 4.0
        act_layer (Type[nn.Module], optional): 使用的激活层类型。默认值 nn.GELU
        norm_layer (Type[nn.Module], optional): 使用的归一化层类型。默认值 nn.BatchNorm2d
        norm_layer_transformer (Type[nn.Module], optional): Transformer 中的归一化层。默认值 nn.LayerNorm
    """

    def __init__(
            self,
            depth: int,
            in_channels: int,
            skip_channels: int,
            out_channels: int,
            num_heads: int = 32,
            grid_window_size: Tuple[int, int] = (7, 7),
            attn_drop: float = 0.,
            drop: float = 0.,
            drop_path: Union[List[float], float] = 0.,
            mlp_ratio: float = 4.,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.BatchNorm2d,
            norm_layer_transformer: Type[nn.Module] = nn.LayerNorm,
            debug=False,
    ) -> None:
        """ 构造函数 """
        # 调用父类构造函数
        super(MaxViTStage, self).__init__()
        self.module_name = '[MaxViTStage]'
        self.debug = debug
        if self.debug: print(self.module_name, '[depth]', depth)
        if self.debug: print(self.module_name, '[in_channels]', in_channels)
        if self.debug: print(self.module_name, '[out_channels]', out_channels)
        if self.debug: print(self.module_name, '[num_heads]', num_heads)
        if self.debug: print(self.module_name, '[grid_window_size]', grid_window_size)
        if self.debug: print(self.module_name, '[attn_drop]', attn_drop)
        if self.debug: print(self.module_name, '[drop]', drop)
        if self.debug: print(self.module_name, '[drop_path]', drop_path)
        if self.debug: print(self.module_name, '[mlp_ratio]', mlp_ratio)
        # 反卷积用于上采样跳跃连接
        self.upsample = DeconvModule(
            in_channels=in_channels,
            out_channels=skip_channels,
            norm_layer=norm_layer,
            act_layer=act_layer)
        # 初始化块
        self.blocks = nn.Sequential(*[
            MaxViTBlock(
                in_channels=2 * skip_channels if index == 0 else out_channels,
                out_channels=out_channels,
                num_heads=num_heads,
                grid_window_size=grid_window_size,
                attn_drop=attn_drop,
                drop=drop,
                drop_path=drop_path if isinstance(drop_path, float) else drop_path[index],
                mlp_ratio=mlp_ratio,
                act_layer=act_layer,
                norm_layer=norm_layer,
                norm_layer_transformer=norm_layer_transformer,
                debug=debug
            )
            for index in range(depth)
        ])

    def forward(self, skip, x):
        """ 前向传播。

        参数:
            skip (torch.Tensor): 输入张量，形状为 [B, C_skip, H, W]。
            x (torch.Tensor): 输入张量，形状为 [B, C_in, H // 2, W // 2]。

        返回:
            output (torch.Tensor): 输出张量，形状为 [B, C_out, H, W]。
        """

        if self.debug: print(self.module_name + '[forward]', f'[x]', x.shape)
        x = self.upsample(x)
        if self.debug: print(self.module_name + '[forward]', f'[upsample][x]', x.shape)
        if self.debug: print(self.module_name + '[forward]', f'[skip]', skip.shape)
        output = torch.cat([skip, x], dim=1)
        output = self.blocks(output)

        return output


class MaxViTDecoder(nn.Module):
    """
    基于 MaxViTBlock 的混合 Transformer 解码器

    参数:
        in_channels (List[int], optional): 每个编码器阶段的输入通道数。默认值 [64, 128, 256, 512]
        depths (Tuple[int, ...], optional): 每个解码器阶段的深度。默认值 (2, 2, 2)
        output_size (Tuple[int, int], optional): 上采样后的输出尺寸。默认值 (256, 256)
        num_heads (int, optional): 注意力头的数量。默认值 32
        grid_window_size (Tuple[int, int], optional): 网格/窗口大小。默认值 (8, 8)
        attn_drop (float, optional): 注意力权重的 Dropout 比例。默认值 0.2
        drop (float, optional): 输出的 Dropout 比例。默认值 0.2
        drop_path (float, optional): 路径的 Dropout 比例。默认值 0.2
        mlp_ratio (float, optional): MLP 隐藏层维度与嵌入维度的比例。默认值 4.0
        act_layer (Type[nn.Module], optional): 使用的激活层。默认值 nn.GELU
        norm_layer (Type[nn.Module], optional): 使用的归一化层。默认值 nn.BatchNorm2d
        norm_layer_transformer (Type[nn.Module], optional): Transformer 中的归一化层。默认值 nn.LayerNorm
        debug (bool, optional): 是否启用调试模式。默认值 False
        num_classes (int, optional): 分割类别数。默认值 21
    """

    def __init__(self, in_channels=[64, 128, 256, 512], depths=(2, 2, 2), output_size=(256, 256),
                 num_heads=32, grid_window_size=(8, 8), attn_drop=0.2, drop=0.2, drop_path=0.2,
                 mlp_ratio=4., act_layer=nn.GELU, norm_layer=nn.BatchNorm2d,
                 norm_layer_transformer=nn.LayerNorm, debug=False, num_classes=21, **kwargs):
        super(MaxViTDecoder, self).__init__()
        self.in_channels = in_channels
        self.conv_seg = nn.Conv2d(in_channels[0], num_classes, kernel_size=1)
        self.dropout = nn.Dropout2d(drop)

        self.module_name = '[MaxViTDecoder]'
        self.debug = debug

        # 初始化阶段
        drop_path_rates = [x.item() for x in torch.linspace(0, drop_path, sum(depths))]
        encoder_stages = len(self.in_channels)
        stages = []
        dp_index = 0
        for index, depth in enumerate(depths):
            dpr = drop_path_rates[dp_index:dp_index + depth]
            dp_index += depth
            stages.append(
                MaxViTStage(
                    depth=depth,
                    in_channels=self.in_channels[encoder_stages - index - 1],
                    skip_channels=self.in_channels[encoder_stages - index - 2] if index < encoder_stages - 1 else
                    self.in_channels[0],
                    out_channels=self.in_channels[encoder_stages - index - 2] if index < encoder_stages - 1 else
                    self.in_channels[0],
                    num_heads=num_heads,
                    grid_window_size=grid_window_size,
                    attn_drop=attn_drop,
                    drop=drop,
                    drop_path=dpr,
                    mlp_ratio=mlp_ratio,
                    act_layer=act_layer,
                    norm_layer=norm_layer,
                    norm_layer_transformer=norm_layer_transformer,
                    debug=debug,
                )
            )
        self.stages = nn.ModuleList(stages)
        self.upsize = nn.Upsample(size=output_size, mode='bilinear', align_corners=False)

    def cls_seg(self, feat):
        """对每个像素进行分类。"""
        if self.dropout is not None:
            feat = self.dropout(feat)
        output = self.conv_seg(feat)
        return output

    @torch.jit.ignore
    def no_weight_decay(self) -> Set[str]:
        """ 获取不应用权重衰减的参数名称。

        返回:
            nwd (Set[str]): 不应用权重衰减的参数名称集合。
        """
        nwd = set()
        for n, _ in self.named_parameters():
            if "relative_position_bias_table" in n:
                nwd.add(n)
        return nwd

    def _forward_feature(self, inputs):
        """ 特征提取的前向传播。

        参数:
            inputs (List[torch.Tensor]): 编码器的输出张量列表，对应不同尺度的特征图。

        返回:
            output (torch.Tensor): 解码器层的输出特征。
        """
        encoder_stages = len(inputs)
        x = inputs[-1]
        if self.debug: print('[MaxViT][forward_features]', f'[bottleneck][x]', x.shape)
        for i in range(len(self.stages)):
            x = self.stages[i](inputs[encoder_stages - i - 2], x)
            if self.debug: print('[MaxViT][forward_features]', f'[stage{i}][x]', x.shape)
        output = self.upsize(x)
        return output

    def forward(self, inputs):
        """前向函数。"""
        TAG = self.module_name + '[forward]'
        if self.debug: print(TAG, '[inputs]', len(inputs))
        output = self._forward_feature(inputs)
        if self.debug: print(TAG, '[output]', output.shape)
        output = self.cls_seg(output)
        return output


if __name__ == '__main__':
    model = MaxViTDecoder(
        in_channels=[64, 128, 256, 512],
        depths=(2, 2, 2),
        grid_window_size=(8, 8),
        attn_drop=0.2,
        drop=0.2,
        drop_path=0.2,
        num_classes=9,  # 根据您的任务设置类别数量
        debug=False,
    )
    # print(model)
    print(f'总参数量: {sum(param.numel() for param in model.parameters()) / 1e6:.2f} M')
    inputs = [
        torch.rand(1, 64, 64, 64),
        torch.rand(1, 128, 32, 32),
        torch.rand(1, 256, 16, 16),
        torch.rand(1, 512, 8, 8)
    ]
    output = model(inputs)
    print(output.shape)
